#include "rerequestor.h"

std::vector< std::pair< uint32_t, int > > Rerequestor::repeatDestinations_;

Rerequestor::Rerequestor(): newSource_(0), newSourcePort_(0), newDestination_(0), newDestinationPort_(0)
{
}

Rerequestor::Rerequestor( const Rerequestor &r )
{
  injector_ = r.injector_;
  newSource_ = r.newSource_;
  newSourcePort_ = r.newSourcePort_;
  newDestination_ = r.newDestination_;
  newDestinationPort_ = r.newDestinationPort_;
}

void Rerequestor::setNewDestination(const uint32_t &ip)
{
  removeIPv4FromRepeatList(newDestination_);
  newDestination_ = ip;
  addIPv4ToRepeatList(newDestination_);
}

void Rerequestor::setNewSource(const uint32_t &ip)
{
  newSource_ = ip;
}

void Rerequestor::removeIPv4FromRepeatList(const uint32_t &ip)
{
  if( ip ) //if its zero then we weren't using it
  {
    std::vector< std::pair<uint32_t, int> > newRepeats;
    for( uint32_t i = 0; i < repeatDestinations_.size(); ++i )
    {
      std::pair< uint32_t, int > set = repeatDestinations_[i];
      if( set.first == ip )
      {
        set.second --;
        if( set.second > 1 )
          newRepeats.push_back( set );
      }
      newRepeats.push_back(set);
    }
    repeatDestinations_ = newRepeats;
  }
}

void Rerequestor::addIPv4ToRepeatList(const uint32_t &ip)
{
  if( ip )
  {
    //see if it is in the list
    bool inList = false;
    for( uint32_t i = 0; i < repeatDestinations_.size(); ++i )
    {
      std::pair< uint32_t, int > set = repeatDestinations_[i];
      if( set.first == ip )
      {
        inList = true;
        repeatDestinations_[i].second += 1;
      }
    }
    if( !inList )
    {
      std::pair< uint32_t, int > newSet;
      newSet.first = ip;
      newSet.second = 1;
      repeatDestinations_.push_back(newSet);
    }
  }
}


bool Rerequestor::isInIPv4RepeatList(const uint32_t &ip)
{
  for( uint32_t i = 0; i < repeatDestinations_.size(); ++i )
  {
    std::pair< uint32_t, int > set = repeatDestinations_[i];
    if( set.first == ip )
    {
      return true;
    }
  }
  return false;
}

void Rerequestor::setNewSourcePort(const uint16_t &port)
{
  newSourcePort_ = port;
}

void Rerequestor::setNewDestinationPort(const uint16_t &port)
{
  newDestinationPort_ = port;
}

uint32_t Rerequestor::newSource() const
{
  return newSource_;
}

uint32_t Rerequestor::newDestination() const
{
  return newDestination_;
}

uint16_t Rerequestor::newSourcePort() const
{
  return newSourcePort_;
}

uint16_t Rerequestor::newDestinationPort() const
{
  return newDestinationPort_;
}

void Rerequestor::processPacket(const Packet &p)
{
  if( p.appSize() && p.appIs< DNS >() )
  {
    DNS dns = p.getApp<DNS>();
    if( dns.isQuery() )
    {      
      Packet newPacket;
      Ethernet eth = p.getLink<Ethernet>();
      uint8_t smac [] = {0x6C, 0xF0, 0x49, 0xED, 0x48, 0xBD};
      eth.setSourceMAC( MACAddress( smac ) );
      newPacket.pushBackLink( eth );
      IPv4 ip;
      if( p.inetSize() && p.inetIs< IPv4 >() )
      {
        ip = p.getInet<IPv4>();
        //check to make sure it wasn't already injected
        if( isInIPv4RepeatList( ip.destinationAddress() ) )
          return;
        if( newSource_ )
          ip.setSourceAddress(newSource_);
        if( newDestination_ )
          ip.setDestinationAddress(newDestination_);
        ip.calculateChecksum();
        newPacket.pushBackInet( ip );
      }
      if( p.transSize() && p.transIs< UDP >() )
      {
        UDP udp = p.getTrans<UDP>();
        if( newSourcePort_ )
          udp.setSourcePort(newSourcePort_);
        if( newDestinationPort_ )
          udp.setDestinationPort(newDestinationPort_);
        udp.calculateChecksum( ip.sourceAddress(), ip.destinationAddress(), dns.makePacket().vector() );
        newPacket.pushBackTrans( udp );
      }
      else if( p.transSize() && p.transIs<TCP>() )
      {
        TCP tcp = p.getTrans<TCP>();
        if( newSourcePort_ )
          tcp.setSourcePort( newSourcePort_ );
        if( newDestinationPort_ )
          tcp.setDestinationPort( newDestinationPort_ );
        newPacket.pushBackTrans( tcp );
      }
      newPacket.pushBackApp( dns );
      injector_.inject(newPacket);
      std::cout << "Packet Away" << std::endl;
    }
  }


}

void Rerequestor::setDevice(const std::string &device)
{
  injector_.setDevice(device);
}

std::string Rerequestor::device() const
{
  return injector_.device();
}

std::string Rerequestor::name() const
{
  return name_;
}

void Rerequestor::setName( const std::string &name  )
{
  name_ = name;
}
